package com.yxsk.relay.job.admin.aop;

import com.alibaba.fastjson.JSONObject;
import com.yxsk.relay.job.component.common.exception.RelayJobRuntimeException;
import com.yxsk.relay.job.component.common.utils.CompressUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.MethodParameter;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.servlet.mvc.method.annotation.RequestBodyAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.Type;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;

/**
 * @Description
 * @Author 11376
 * @CreateTime 2019/8/24 11:40
 */
@Slf4j
public class CompressMessageAspectAware implements RequestBodyAdvice, ResponseBodyAdvice<Object> {

    private static final int BUFFER_SIZE = 128;

    private static final Charset MESSAGE_CHARSET = StandardCharsets.UTF_8;

    @Override
    public boolean supports(MethodParameter methodParameter, Type type, Class<? extends HttpMessageConverter<?>> aClass) {
        return true;
    }

    @Override
    public HttpInputMessage beforeBodyRead(HttpInputMessage httpInputMessage, MethodParameter methodParameter, Type type, Class<? extends HttpMessageConverter<?>> aClass) throws IOException {
        try (InputStream is = httpInputMessage.getBody()) {
            // 缓存流
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            byte[] buffer = new byte[BUFFER_SIZE];
            int offset = -1;
            while ((offset = is.read(buffer, 0, BUFFER_SIZE)) != -1) {
                bos.write(buffer, 0, offset);
            }

            byte[] httpBytes = bos.toByteArray();

            // 解压数据包
            byte[] message = CompressUtils.uncompress(httpBytes, CompressUtils.CompressType.GZIP);

            // 获取 HTTP 请求数据
            if (log.isDebugEnabled()) {
                log.debug("HTTP input message: {}", httpBytes == null ? "" : new String(message, MESSAGE_CHARSET));
            }

            return new RelayJobHttpMessage(new ByteArrayInputStream(message), httpInputMessage.getHeaders());
        }
    }

    @Override
    public Object afterBodyRead(Object body, HttpInputMessage httpInputMessage, MethodParameter methodParameter, Type type, Class<? extends HttpMessageConverter<?>> aClass) {
        return body;
    }

    @Override
    public Object handleEmptyBody(Object body, HttpInputMessage httpInputMessage, MethodParameter methodParameter, Type type, Class<? extends HttpMessageConverter<?>> aClass) {
        return body;
    }


    @Override
    public boolean supports(MethodParameter methodParameter, Class<? extends HttpMessageConverter<?>> aClass) {
        return true;
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter methodParameter, MediaType mediaType, Class<? extends HttpMessageConverter<?>> aClass, ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse) {
        if (body == null) {
            if (log.isDebugEnabled()) {
                log.debug("HTTP response message is empty.");
            }
            return null;
        }

        if (mediaType.includes(MediaType.APPLICATION_JSON)) {
            // body 转 json
            String json = JSONObject.toJSONString(body);
            if (log.isDebugEnabled()) {
                log.debug("HTTP response message: {}", json);
            }

            // 数据压缩
            try {
                byte[] message = CompressUtils.compress(json.getBytes(MESSAGE_CHARSET), CompressUtils.CompressType.GZIP);
                serverHttpResponse.getBody().write(message);
                return null;
            } catch (IOException e) {
                log.error("Message compress error", e);
                throw new RelayJobRuntimeException(e);
            }
        }

        // 无处理
        return body;
    }
}
